from langchain_openai import ChatOpenAI
from langchain_core.prompts import ChatPromptTemplate
from langgraph.graph import END, StateGraph

from config import settings
from data_model.data_schemas import UserInfo, ExtractionState
from src.utils import empty_user_info_detect


class ExtractorAgent:
    def __init__(self):
        self.llm = ChatOpenAI(
                    model=settings.MODEL_NAME,
                    base_url=settings.OLLAMA_URL,
                    api_key=settings.OLLAMA_API_KEY
                )
        self.prompt = ChatPromptTemplate.from_messages(
                        [
                            (
                                "system",
                                """
                                提取用户输入中的生日，星座以及问题，结构化输出为 'json'。\n
                                ##注意##\n
                                1.若其中一项或多项为空，则你提取为空值。\n
                                """,
                            ),
                            ("human", "{query}"),
                        ]
                    )
        self.workflow = self._build_workflow()

    def _extraction_node(self, state:ExtractionState)->dict:
        structured_llm = self.llm.with_structured_output(UserInfo)
        chain = self.prompt | structured_llm
        response = chain.invoke({"query":state.query})
        return {"answer": response}

    @staticmethod
    def _requests_missing_info_node(state:ExtractionState):
        # 处理信息不完整的情况，返回具体信息
        missing_info_msg = empty_user_info_detect(state.answer)
        return {"answer": missing_info_msg}

    def routing_node(self,state:ExtractionState)->bool:
        # 添加条件判断节点，若用户信息完整，则返回True；否则返回False。
        is_complete = empty_user_info_detect(state.answer)
        if is_complete is True:
            return "complete"
        else:
            return "incomplete"

    def _build_workflow(self):
        # 创建工作流
        workflow = StateGraph(ExtractionState)
        workflow.add_node('extraction', self._extraction_node)
        workflow.add_node('missing_info_node', self._requests_missing_info_node)
        # 添加入口点
        workflow.set_entry_point('extraction')

        # 添加条件边
        workflow.add_conditional_edges("extraction",self.routing_node,
                                       {"complete":END, "incomplete":"missing_info_node"})

        workflow.add_edge('missing_info_node',END)

        return workflow.compile()

    def chat(self, query:str)->UserInfo:
        initial_state = ExtractionState(query=query)
        result = self.workflow.invoke(initial_state)
        return result["answer"]

# # 使用工作流
# if __name__ == "__main__":
#     agent = ExtractorAgent()
#     resp = agent.chat("你好")
#     print(resp)
